package common

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"log"
	"os"
)

// 生成rsa密钥对，并且保存到conf目录下
func GenerateRsaKey(keySize int, pubPath, priPath *string) {
	// 使用rsa中的GenertaeKey生成私钥
	privateKey, err := rsa.GenerateKey(rand.Reader, keySize)
	if err != nil {
		panic(err)
	}

	// 通过x509标准将得到的rsa私钥序列化为ASN.1的DER编码格式
	derText := x509.MarshalPKCS1PrivateKey(privateKey)

	// 要组织一个pem.block
	block := pem.Block{
		Type:  "rsa private key", //这个地方随便写什么
		Bytes: derText,           //ASN.1的DER编码格式
	}
	// pem编码
	file, err1 := os.Create(*priPath) //保存在磁盘的文件名
	if err1 != nil {
		panic(err1)
	}
	pem.Encode(file, &block)
	file.Close()

	// 从私钥中取出公钥
	publicKey := privateKey.PublicKey
	// 使用x509标准序列化  注意参数传参为地址
	derstream, err2 := x509.MarshalPKIXPublicKey(&publicKey)
	if err2 != nil {
		panic(err2)
	}
	// 将得到的数据放到pem.block中
	block = pem.Block{
		Type:    "rsa public key",
		Headers: nil,
		Bytes:   derstream,
	}
	// pem编码
	file, err = os.Create(*pubPath)
	if err != nil {
		panic(err)
	}
	pem.Encode(file, &block)
	file.Close()
}

/**
 * rsa加密
 */
func RsaEncrypt(origData []byte, publicKey []byte) (string, error) {
	defer func() {
		err := recover()
		if err != nil {
			log.Println(string(origData), string(publicKey), "encrypt failed:", err)
		}
	}()

	block, _ := pem.Decode(publicKey)
	if block == nil {
		return "", errors.New("public key error")
	}
	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return "", err
	}
	pub := pubInterface.(*rsa.PublicKey)

	partLen := pub.N.BitLen()/8 - 11
	chunks := rsa_split(origData, partLen)

	buffer := bytes.NewBufferString("")
	for _, chunk := range chunks {
		bytes, err := rsa.EncryptPKCS1v15(rand.Reader, pub, chunk)
		if err != nil {
			return "", err
		}
		buffer.Write(bytes)
	}
	return base64.StdEncoding.EncodeToString(buffer.Bytes()), nil
}

/**
 * rsa解密
 */
func RsaDecrypt(ciphertext string, publicKey []byte, privateKey []byte) (string, error) {
	defer func() {
		err := recover()
		if err != nil {
			log.Println(ciphertext, string(publicKey), string(privateKey), "encrypt failed:", err)
		}
	}()

	pub_block, _ := pem.Decode(publicKey)
	if pub_block == nil {
		return "", errors.New("public key error")
	}
	pubInterface, err := x509.ParsePKIXPublicKey(pub_block.Bytes)
	if err != nil {
		return "", err
	}
	pub := pubInterface.(*rsa.PublicKey)

	block, _ := pem.Decode(privateKey)
	if block == nil {
		return "", errors.New("private key error!")
	}
	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}

	partLen := pub.N.BitLen() / 8
	raw, err := base64.StdEncoding.DecodeString(ciphertext)
	chunks := rsa_split([]byte(raw), partLen) // 将字节数据分段返回

	buffer := bytes.NewBufferString("")
	for _, chunk := range chunks {
		decrypted, err := rsa.DecryptPKCS1v15(rand.Reader, priv, chunk)
		if err != nil {
			return "", err
		}
		buffer.Write(decrypted)
	}
	return buffer.String(), err
}

/**
 * rsa分段处理
 */
func rsa_split(buf []byte, lim int) [][]byte {
	var chunk []byte
	chunks := make([][]byte, 0, len(buf)/lim+1)
	for len(buf) >= lim {
		chunk, buf = buf[:lim], buf[lim:]
		chunks = append(chunks, chunk)
	}
	if len(buf) > 0 {
		chunks = append(chunks, buf[:len(buf)])
	}
	return chunks
}
